use rmcp::{
    model::{ClientCapabilities, ClientInfo, Implementation, Request, Tool},
    transport::StreamableHttpClientTransport,
    ServiceExt,
};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fs;
use std::path::Path;
use std::process::Stdio;
use std::sync::Arc;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::process::Command;
use tokio::sync::{Mutex, RwLock};

#[derive(Debug, Clone)]
pub enum TransportType {
    Stdio,
    Http,
    Sse,
}

#[derive(Debug, Clone)]
pub enum TransportConfig {
    Stdio {
        command: String,
        args: Vec<String>,
        env_vars: std::collections::HashMap<String, String>,
    },
    Http {
        url: String,
        headers: std::collections::HashMap<String, String>,
    },
    Sse {
        url: String,
        headers: std::collections::HashMap<String, String>,
    },
}

#[derive(Debug, Clone)]
pub struct ServiceId(String);

impl ServiceId {
    pub fn value(&self) -> &str {
        &self.0
    }
    pub fn new(name: String) -> Self {
        ServiceId(name)
    }
}

#[derive(Debug, Clone)]
pub enum ServiceType {
    Mcp,
}

#[derive(Debug, Clone)]
pub struct ServiceConfig {
    pub id: ServiceId,
    pub name: String,
    pub service_type: ServiceType,
    pub transport: TransportConfig,
    pub enabled: bool,
}

// Simple response wrapper
#[derive(Debug)]
pub struct SimpleResponse {
    pub is_error: bool,
    pub error_message: Option<String>,
    pub result_value: Option<serde_json::Value>,
}

impl SimpleResponse {
    pub fn is_error(&self) -> bool {
        self.is_error
    }

    pub fn error(&self) -> Option<SimpleError> {
        self.error_message.as_ref().map(|msg| SimpleError {
            message: msg.clone(),
        })
    }

    pub fn result(&self) -> Option<serde_json::Value> {
        self.result_value.clone()
    }
}

#[derive(Debug)]
pub struct SimpleError {
    pub message: String,
}

// Custom stdio transport for MCP
#[derive(Debug)]
pub struct StdioTransport {
    child: Arc<Mutex<tokio::process::Child>>,
    stdin: Arc<Mutex<tokio::process::ChildStdin>>,
    stdout: Arc<Mutex<BufReader<tokio::process::ChildStdout>>>,
}

impl StdioTransport {
    pub async fn new(
        command: &str,
        args: &[String],
        env_vars: &HashMap<String, String>,
    ) -> anyhow::Result<Self> {
        let mut cmd = Command::new(command);
        cmd.args(args)
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped());

        // Add environment variables
        for (key, value) in env_vars {
            cmd.env(key, value);
        }

        let mut child = cmd
            .spawn()
            .map_err(|e| anyhow::anyhow!("Failed to spawn process '{}': {}", command, e))?;

        let stdin = child
            .stdin
            .take()
            .ok_or_else(|| anyhow::anyhow!("Failed to get stdin handle for process"))?;

        let stdout = child
            .stdout
            .take()
            .ok_or_else(|| anyhow::anyhow!("Failed to get stdout handle for process"))?;

        Ok(Self {
            child: Arc::new(Mutex::new(child)),
            stdin: Arc::new(Mutex::new(stdin)),
            stdout: Arc::new(Mutex::new(BufReader::new(stdout))),
        })
    }

    pub async fn send_request(&self, request: &str) -> anyhow::Result<String> {
        // Check if process is still alive
        if !self.is_alive().await {
            return Err(anyhow::anyhow!("Process is not running"));
        }

        // Send request to stdin
        {
            let mut stdin = self.stdin.lock().await;
            stdin
                .write_all(request.as_bytes())
                .await
                .map_err(|e| anyhow::anyhow!("Failed to write to stdin: {}", e))?;
            stdin
                .write_all(b"\n")
                .await
                .map_err(|e| anyhow::anyhow!("Failed to write newline to stdin: {}", e))?;
            stdin
                .flush()
                .await
                .map_err(|e| anyhow::anyhow!("Failed to flush stdin: {}", e))?;
        }

        // Read response from stdout with timeout
        let mut stdout = self.stdout.lock().await;
        let mut response = String::new();

        // Use timeout to avoid hanging indefinitely
        match tokio::time::timeout(
            std::time::Duration::from_secs(30),
            stdout.read_line(&mut response),
        )
        .await
        {
            Ok(Ok(bytes_read)) => {
                if bytes_read == 0 {
                    return Err(anyhow::anyhow!("Process closed stdout"));
                }

                let trimmed = response.trim();

                // Try to parse as JSON to validate completeness
                if let Err(_) = serde_json::from_str::<serde_json::Value>(trimmed) {
                    // If JSON is incomplete, try to read more lines
                    let mut additional_lines = String::new();
                    let mut attempts = 0;
                    const MAX_ATTEMPTS: usize = 10;

                    while attempts < MAX_ATTEMPTS {
                        match tokio::time::timeout(
                            std::time::Duration::from_millis(100),
                            stdout.read_line(&mut additional_lines),
                        )
                        .await
                        {
                            Ok(Ok(0)) => break, // EOF
                            Ok(Ok(_)) => {
                                response.push_str(&additional_lines);
                                additional_lines.clear();

                                // Try parsing again
                                if serde_json::from_str::<serde_json::Value>(response.trim())
                                    .is_ok()
                                {
                                    break;
                                }
                            }
                            Ok(Err(_)) | Err(_) => break,
                        }
                        attempts += 1;
                    }
                }

                Ok(response.trim().to_string())
            }
            Ok(Err(e)) => Err(anyhow::anyhow!("Failed to read from stdout: {}", e)),
            Err(_) => Err(anyhow::anyhow!("Timeout waiting for response")),
        }
    }

    pub async fn is_alive(&self) -> bool {
        let mut child = self.child.lock().await;
        match child.try_wait() {
            Ok(Some(_)) => false, // Process has exited
            Ok(None) => true,     // Process is still running
            Err(_) => false,      // Error checking status
        }
    }

    pub async fn terminate(&self) -> anyhow::Result<()> {
        let mut child = self.child.lock().await;
        if let Err(e) = child.kill().await {
            eprintln!("Failed to kill child process: {}", e);
        }
        let _ = child.wait().await;
        Ok(())
    }
}

// Service wrapper for MCP connections
#[derive(Debug, Clone)]
pub struct ServiceWrapper {
    inner: Arc<ServiceWrapperInner>,
}

#[derive(Debug)]
struct ServiceWrapperInner {
    client: Arc<Mutex<Box<dyn std::any::Any + Send + Sync>>>,
    stdio_transport: Option<Arc<StdioTransport>>,
    name: String,
    transport_type: TransportType,
}

impl ServiceWrapper {
    pub async fn new(config: ServiceConfig) -> anyhow::Result<Self> {
        // Create transport based on config type
        match &config.transport {
            TransportConfig::Stdio {
                command,
                args,
                env_vars,
            } => {
                // Create stdio transport
                let stdio_transport = StdioTransport::new(command, args, env_vars).await?;

                println!("Initializing MCP connection for service: {}", config.name);

                // Initialize MCP connection over stdio
                let init_request = serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 1,
                    "method": "initialize",
                    "params": {
                        "protocolVersion": "2024-11-05",
                        "capabilities": {
                            "roots": {
                                "listChanged": true
                            },
                            "sampling": {}
                        },
                        "clientInfo": {
                            "name": "rig-app",
                            "version": "0.1.0"
                        }
                    }
                });

                let response = stdio_transport
                    .send_request(&init_request.to_string())
                    .await
                    .map_err(|e| anyhow::anyhow!("Failed to initialize MCP connection: {}", e))?;

                println!("MCP initialization response: {}", response);

                // Parse initialization response
                if let Ok(response_json) = serde_json::from_str::<serde_json::Value>(&response) {
                    if let Some(error) = response_json.get("error") {
                        return Err(anyhow::anyhow!("MCP initialization failed: {}", error));
                    }

                    // Log server capabilities if available
                    if let Some(result) = response_json.get("result") {
                        if let Some(capabilities) = result.get("capabilities") {
                            println!("Server capabilities: {}", capabilities);
                        }
                    }
                }

                // Send initialized notification
                let initialized_notification = serde_json::json!({
                    "jsonrpc": "2.0",
                    "method": "notifications/initialized"
                });

                if let Err(e) = stdio_transport
                    .send_request(&initialized_notification.to_string())
                    .await
                {
                    eprintln!("Warning: Failed to send initialized notification: {}", e);
                }

                println!("MCP service '{}' initialized successfully", config.name);

                Ok(Self {
                    inner: Arc::new(ServiceWrapperInner {
                        client: Arc::new(Mutex::new(Box::new(()))), // Placeholder for stdio
                        stdio_transport: Some(Arc::new(stdio_transport)),
                        name: config.name,
                        transport_type: TransportType::Stdio,
                    }),
                })
            }
            TransportConfig::Http { url, headers: _ }
            | TransportConfig::Sse { url, headers: _ } => {
                // Create HTTP/SSE transport
                let transport = StreamableHttpClientTransport::from_uri(url.as_str());

                let client_info = ClientInfo {
                    protocol_version: Default::default(),
                    capabilities: ClientCapabilities::default(),
                    client_info: Implementation {
                        name: "rig-app".to_string(),
                        version: "0.1.0".to_string(),
                    },
                };

                let client = client_info
                    .serve(transport)
                    .await
                    .map_err(|e| anyhow::anyhow!("Failed to connect to MCP server: {}", e))?;

                let transport_type = match config.transport {
                    TransportConfig::Http { .. } => TransportType::Http,
                    TransportConfig::Sse { .. } => TransportType::Sse,
                    _ => unreachable!(),
                };

                Ok(Self {
                    inner: Arc::new(ServiceWrapperInner {
                        client: Arc::new(Mutex::new(Box::new(client))),
                        stdio_transport: None,
                        name: config.name,
                        transport_type,
                    }),
                })
            }
        }
    }

    pub async fn is_healthy(&self) -> anyhow::Result<bool> {
        match self.inner.transport_type {
            TransportType::Stdio => {
                if let Some(stdio_transport) = &self.inner.stdio_transport {
                    Ok(stdio_transport.is_alive().await)
                } else {
                    Ok(false)
                }
            }
            TransportType::Http | TransportType::Sse => {
                // For HTTP/SSE, we assume healthy if client exists
                // In a real implementation, you might want to send a ping request
                Ok(true)
            }
        }
    }

    pub async fn send_request(
        &self,
        request: Request<String, serde_json::Value>,
    ) -> anyhow::Result<SimpleResponse> {
        match self.inner.transport_type {
            TransportType::Stdio => self.send_stdio_request(request).await,
            TransportType::Http | TransportType::Sse => self.send_http_request(request).await,
        }
    }

    async fn send_stdio_request(
        &self,
        request: Request<String, serde_json::Value>,
    ) -> anyhow::Result<SimpleResponse> {
        let stdio_transport = self
            .inner
            .stdio_transport
            .as_ref()
            .ok_or_else(|| anyhow::anyhow!("Stdio transport not available"))?;

        // Parse the method to extract operation type
        let parts: Vec<&str> = request.method.split('/').collect();

        if parts.len() < 2 {
            return Err(anyhow::anyhow!("Invalid request method format"));
        }

        match parts[1] {
            "list" => {
                // Send tools/list request
                let list_request = serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 2,
                    "method": "tools/list"
                });

                let response = stdio_transport
                    .send_request(&list_request.to_string())
                    .await?;

                match serde_json::from_str::<serde_json::Value>(&response) {
                    Ok(json_response) => {
                        if let Some(error) = json_response.get("error") {
                            Ok(SimpleResponse {
                                is_error: true,
                                error_message: Some(error.to_string()),
                                result_value: None,
                            })
                        } else if let Some(result) = json_response.get("result") {
                            Ok(SimpleResponse {
                                is_error: false,
                                error_message: None,
                                result_value: Some(result.clone()),
                            })
                        } else {
                            Ok(SimpleResponse {
                                is_error: false,
                                error_message: None,
                                result_value: Some(serde_json::json!({"tools": []})),
                            })
                        }
                    }
                    Err(e) => Ok(SimpleResponse {
                        is_error: true,
                        error_message: Some(format!("Failed to parse response: {}", e)),
                        result_value: None,
                    }),
                }
            }
            "call" => {
                if parts.len() < 3 {
                    return Err(anyhow::anyhow!("Tool name not specified"));
                }
                let tool_name = parts[2];

                // Send tools/call request
                let call_request = serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 3,
                    "method": "tools/call",
                    "params": {
                        "name": tool_name,
                        "arguments": request.params
                    }
                });

                let response = stdio_transport
                    .send_request(&call_request.to_string())
                    .await?;

                match serde_json::from_str::<serde_json::Value>(&response) {
                    Ok(json_response) => {
                        if let Some(error) = json_response.get("error") {
                            Ok(SimpleResponse {
                                is_error: true,
                                error_message: Some(error.to_string()),
                                result_value: None,
                            })
                        } else if let Some(result) = json_response.get("result") {
                            Ok(SimpleResponse {
                                is_error: false,
                                error_message: None,
                                result_value: Some(result.clone()),
                            })
                        } else {
                            Ok(SimpleResponse {
                                is_error: false,
                                error_message: None,
                                result_value: Some(serde_json::json!({
                                    "content": [
                                        {
                                            "type": "text",
                                            "text": "Tool executed successfully"
                                        }
                                    ]
                                })),
                            })
                        }
                    }
                    Err(e) => Ok(SimpleResponse {
                        is_error: true,
                        error_message: Some(format!("Failed to parse response: {}", e)),
                        result_value: None,
                    }),
                }
            }
            _ => Err(anyhow::anyhow!("Unknown request type: {}", parts[1])),
        }
    }

    async fn send_http_request(
        &self,
        request: Request<String, serde_json::Value>,
    ) -> anyhow::Result<SimpleResponse> {
        // Parse the method to extract tool name
        // Expected format: "service_name/call/tool_name" or "service_name/list"
        let parts: Vec<&str> = request.method.split('/').collect();

        if parts.len() < 2 {
            return Err(anyhow::anyhow!("Invalid request method format"));
        }

        match parts[1] {
            "list" => {
                // For HTTP/SSE, return empty list for now
                // In a real implementation, you would use the rmcp client to list tools
                Ok(SimpleResponse {
                    is_error: false,
                    error_message: None,
                    result_value: Some(serde_json::json!({"tools": []})),
                })
            }
            "call" => {
                if parts.len() < 3 {
                    return Err(anyhow::anyhow!("Tool name not specified"));
                }
                let tool_name = parts[2];

                // Convert params to arguments map if it's an object
                let arguments = if let serde_json::Value::Object(map) = request.params {
                    Some(map)
                } else if request.params.is_null() {
                    None
                } else {
                    // Wrap non-object params in a map with "value" key
                    let mut map = serde_json::Map::new();
                    map.insert("value".to_string(), request.params);
                    Some(map)
                };

                // For HTTP/SSE, return placeholder response
                // In a real implementation, you would use the rmcp client to call tools
                Ok(SimpleResponse {
                    is_error: false,
                    error_message: None,
                    result_value: Some(serde_json::json!({
                        "content": [
                            {
                                "type": "text",
                                "text": format!("Tool '{}' executed with arguments: {:?}", tool_name, arguments)
                            }
                        ]
                    })),
                })
            }
            _ => Err(anyhow::anyhow!("Unknown request type: {}", parts[1])),
        }
    }

    pub fn name(&self) -> &str {
        &self.inner.name
    }

    pub async fn cleanup(&self) -> anyhow::Result<()> {
        if let Some(stdio_transport) = &self.inner.stdio_transport {
            stdio_transport.terminate().await?;
        }
        Ok(())
    }
}

// MCP Tool Management
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolDefinition {
    pub name: String,
    pub description: String,
    pub parameters: Option<serde_json::Value>,
}

// MCP Tool execution request and response
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MCPToolCallRequest {
    pub tool_name: String,
    pub arguments: serde_json::Value,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MCPToolCallResponse {
    pub success: bool,
    pub result: Option<serde_json::Value>,
    pub error: Option<String>,
}

// MCP Service Configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum MCPProtocol {
    Stdio,
    Sse,
    Http,
}

impl From<MCPProtocol> for TransportType {
    fn from(protocol: MCPProtocol) -> Self {
        match protocol {
            MCPProtocol::Stdio => TransportType::Stdio,
            MCPProtocol::Sse => TransportType::Sse,
            MCPProtocol::Http => TransportType::Http,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MCPServiceConfig {
    pub id: i64,
    pub name: String,
    pub protocol: MCPProtocol,
    pub command: Option<String>,
    pub args: Option<Vec<String>>,
    pub url: Option<String>,
    pub headers: Option<HashMap<String, String>>,
    pub env_vars: Option<HashMap<String, String>>,
    pub enabled: bool,
    pub created_at: String,
}

impl TryFrom<MCPServiceConfig> for ServiceConfig {
    type Error = anyhow::Error;

    fn try_from(config: MCPServiceConfig) -> Result<Self, Self::Error> {
        let transport_config = match config.protocol {
            MCPProtocol::Stdio => {
                let command = config
                    .command
                    .ok_or_else(|| anyhow::anyhow!("Command required for stdio protocol"))?;
                TransportConfig::Stdio {
                    command,
                    args: config.args.unwrap_or_default(),
                    env_vars: config.env_vars.unwrap_or_default(),
                }
            }
            MCPProtocol::Http => {
                let url = config
                    .url
                    .ok_or_else(|| anyhow::anyhow!("URL required for HTTP protocol"))?;
                TransportConfig::Http {
                    url,
                    headers: config.headers.unwrap_or_default(),
                }
            }
            MCPProtocol::Sse => {
                let url = config
                    .url
                    .ok_or_else(|| anyhow::anyhow!("URL required for SSE protocol"))?;
                TransportConfig::Sse {
                    url,
                    headers: config.headers.unwrap_or_default(),
                }
            }
        };

        Ok(ServiceConfig {
            id: ServiceId::new(config.name.clone()),
            name: config.name,
            service_type: ServiceType::Mcp,
            transport: transport_config,
            enabled: config.enabled,
        })
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MCPToolDefinition {
    pub service_id: i64,
    pub service_name: String,
    pub name: String,
    pub description: String,
    pub input_schema: serde_json::Value,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct MCPData {
    tools: Vec<ToolDefinition>,
    mcp_services: Vec<MCPServiceConfig>,
    mcp_tools: Vec<MCPToolDefinition>,
}

#[derive(Debug, Clone)]
pub struct ToolManager {
    data_path: String,
    data: Arc<Mutex<MCPData>>,
    pub service_registry: Arc<RwLock<HashMap<i64, ServiceWrapper>>>,
}

impl ToolManager {
    pub fn new<P: AsRef<Path>>(data_path: P) -> anyhow::Result<Self> {
        let data_path = data_path.as_ref().to_string_lossy().to_string();
        let manager = Self {
            data_path: data_path.clone(),
            data: Arc::new(Mutex::new(MCPData {
                tools: vec![],
                mcp_services: vec![],
                mcp_tools: vec![],
            })),
            service_registry: Arc::new(RwLock::new(HashMap::new())),
        };
        manager.load_data()?;
        Ok(manager)
    }

    // Initialize all enabled MCP services from data
    pub async fn initialize_services(&self) -> anyhow::Result<()> {
        let services = self.list_mcp_services()?;

        for service in services {
            if service.enabled {
                if let Err(e) = self.register_mcp_service(&service).await {
                    eprintln!("Failed to initialize service {}: {}", service.name, e);
                }
            }
        }

        Ok(())
    }

    fn load_data(&self) -> anyhow::Result<()> {
        if Path::new(&self.data_path).exists() {
            let content = fs::read_to_string(&self.data_path)?;
            let mcp_data: MCPData = serde_json::from_str(&content)?;
            *self.data.blocking_lock() = mcp_data;
        }
        Ok(())
    }

    fn save_data(&self) -> anyhow::Result<()> {
        let content = serde_json::to_string_pretty(&*self.data.blocking_lock())?;
        fs::write(&self.data_path, content)?;
        Ok(())
    }

    fn get_next_service_id(&self) -> i64 {
        let data = self.data.blocking_lock();
        data.mcp_services.iter().map(|s| s.id).max().unwrap_or(0) + 1
    }

    pub async fn register_tool(&self, tool: &ToolDefinition) -> anyhow::Result<()> {
        let mut data = self.data.lock().await;
        data.tools.push(tool.clone());
        drop(data);
        self.save_data()?;
        Ok(())
    }

    pub fn list_tools(&self) -> anyhow::Result<Vec<ToolDefinition>> {
        let data = self.data.blocking_lock();
        Ok(data.tools.clone())
    }

    pub fn get_agent_tools(&self, _agent_id: i64) -> anyhow::Result<Vec<ToolDefinition>> {
        // This method now requires AgentManager to be passed in
        // For now, return empty list - this will be handled differently in the new architecture
        Ok(vec![])
    }

    // MCP Service Management
    pub async fn register_mcp_service(&self, service: &MCPServiceConfig) -> anyhow::Result<i64> {
        let service_id = if service.id > 0 {
            service.id
        } else {
            self.get_next_service_id()
        };

        let mut service_with_id = service.clone();
        service_with_id.id = service_id;

        let mut data = self.data.blocking_lock();
        data.mcp_services.push(service_with_id);
        drop(data);
        self.save_data()?;

        // Create rmcp Service and register it
        if service.enabled {
            let service_config: ServiceConfig = service.clone().try_into()?;
            let rmcp_service = ServiceWrapper::new(service_config).await?;

            let mut registry = self.service_registry.write().await;
            registry.insert(service_id, rmcp_service);

            // Automatically discover and register tools for this service
            if let Err(e) = self.refresh_mcp_tools(service_id).await {
                eprintln!(
                    "Failed to discover tools for service {}: {}",
                    service.name, e
                );
            }
        }

        Ok(service_id)
    }

    pub fn list_mcp_services(&self) -> anyhow::Result<Vec<MCPServiceConfig>> {
        let data = self.data.blocking_lock();
        Ok(data.mcp_services.clone())
    }

    pub fn get_mcp_service(&self, service_id: i64) -> anyhow::Result<Option<MCPServiceConfig>> {
        let data = self.data.blocking_lock();
        Ok(data
            .mcp_services
            .iter()
            .find(|s| s.id == service_id)
            .cloned())
    }

    pub fn delete_mcp_service(&self, service_id: i64) -> anyhow::Result<bool> {
        let mut data = self.data.blocking_lock();
        let initial_len = data.mcp_services.len();
        data.mcp_services.retain(|service| service.id != service_id);
        data.mcp_tools.retain(|tool| tool.service_id != service_id);
        let deleted = data.mcp_services.len() < initial_len;
        drop(data);

        if deleted {
            self.save_data()?;
        }
        Ok(deleted)
    }

    pub fn register_mcp_tools(
        &self,
        service_id: i64,
        tools: &[MCPToolDefinition],
    ) -> anyhow::Result<()> {
        let mut data = self.data.blocking_lock();

        // Remove existing tools for this service
        data.mcp_tools.retain(|tool| tool.service_id != service_id);

        // Add new tools
        for tool in tools {
            data.mcp_tools.push(tool.clone());
        }
        drop(data);

        self.save_data()?;
        Ok(())
    }

    pub fn list_mcp_tools(
        &self,
        service_id: Option<i64>,
    ) -> anyhow::Result<Vec<MCPToolDefinition>> {
        let data = self.data.blocking_lock();
        if let Some(id) = service_id {
            Ok(data
                .mcp_tools
                .iter()
                .filter(|tool| tool.service_id == id)
                .cloned()
                .collect())
        } else {
            Ok(data.mcp_tools.clone())
        }
    }

    pub fn get_all_tools(&self) -> anyhow::Result<Vec<ToolDefinition>> {
        let mut all_tools = self.list_tools()?;

        // Add MCP tools
        let mcp_tools = self.list_mcp_tools(None)?;
        for mcp_tool in mcp_tools {
            all_tools.push(ToolDefinition {
                name: format!("{}::{}", mcp_tool.service_name, mcp_tool.name),
                description: mcp_tool.description,
                parameters: Some(mcp_tool.input_schema),
            });
        }

        Ok(all_tools)
    }

    pub fn get_agent_all_tools(&self, _agent_id: i64) -> anyhow::Result<Vec<ToolDefinition>> {
        // This method now requires AgentManager to be passed in
        // For now, return all tools - this will be handled differently in the new architecture
        self.get_all_tools()
    }

    // MCP Tool execution
    pub async fn execute_mcp_tool(
        &self,
        service_id: i64,
        tool_name: &str,
        arguments: serde_json::Value,
    ) -> anyhow::Result<MCPToolCallResponse> {
        let service_config = self
            .get_mcp_service(service_id)?
            .ok_or_else(|| anyhow::anyhow!("Service config not found for id: {}", service_id))?;

        let registry = self.service_registry.read().await;
        let service = registry
            .get(&service_id)
            .ok_or_else(|| anyhow::anyhow!("Service not found: {}", service_id))?;

        // Check service health before executing
        match service.is_healthy().await {
            Ok(true) => {}
            Ok(false) => {
                return Ok(MCPToolCallResponse {
                    success: false,
                    result: None,
                    error: Some("Service is unhealthy".to_string()),
                });
            }
            Err(e) => {
                return Ok(MCPToolCallResponse {
                    success: false,
                    result: None,
                    error: Some(format!("Service health check failed: {}", e)),
                });
            }
        }

        let request = Request {
            method: format!("{}/call/{}", service_config.name, tool_name),
            params: arguments,
            extensions: Default::default(),
        };

        match service.send_request(request).await {
            Ok(response) => {
                if response.is_error() {
                    Ok(MCPToolCallResponse {
                        success: false,
                        result: None,
                        error: response.error().map(|e| e.message.clone()),
                    })
                } else {
                    Ok(MCPToolCallResponse {
                        success: true,
                        result: response.result(),
                        error: None,
                    })
                }
            }
            Err(e) => Ok(MCPToolCallResponse {
                success: false,
                result: None,
                error: Some(format!("Failed to send request: {}", e)),
            }),
        }
    }

    // Find MCP tool by name (format: "service_name::tool_name")
    pub fn find_mcp_tool(
        &self,
        tool_name: &str,
    ) -> anyhow::Result<Option<(i64, MCPToolDefinition)>> {
        if let Some((service_name, tool_name)) = tool_name.split_once("::") {
            let data = self.data.blocking_lock();

            // Find service by name first
            let service = data.mcp_services.iter().find(|s| s.name == service_name);
            if let Some(service) = service {
                let service_id = service.id;
                // Find tool by service_id and tool_name
                let tool = data
                    .mcp_tools
                    .iter()
                    .find(|t| t.service_id == service_id && t.name == tool_name);
                if let Some(tool) = tool {
                    return Ok(Some((service_id, tool.clone())));
                }
            }
        }
        Ok(None)
    }

    // Discover tools from MCP service using rmcp
    pub async fn discover_mcp_tools(
        &self,
        service_id: i64,
    ) -> anyhow::Result<Vec<MCPToolDefinition>> {
        let service_config = self
            .get_mcp_service(service_id)?
            .ok_or_else(|| anyhow::anyhow!("Service config not found for id: {}", service_id))?;

        let registry = self.service_registry.read().await;
        let service = registry
            .get(&service_id)
            .ok_or_else(|| anyhow::anyhow!("Service not found: {}", service_id))?;

        let request = Request {
            method: format!("{}/list", service.name()),
            params: serde_json::Value::Null,
            extensions: Default::default(),
        };

        let response = service.send_request(request).await?;

        if response.is_error() {
            return Err(anyhow::anyhow!(
                "Failed to discover tools: {:?}",
                response.error()
            ));
        }

        if let Some(result) = response.result() {
            // Handle different response formats
            let tools_array = if let Some(tools) = result.get("tools") {
                tools.clone()
            } else {
                result.clone()
            };

            if let Ok(discovered_tools) =
                serde_json::from_value::<Vec<serde_json::Value>>(tools_array)
            {
                let mut tools = Vec::new();

                for tool_value in discovered_tools {
                    if let Ok(tool) = serde_json::from_value::<Tool>(tool_value.clone()) {
                        // Use rmcp Tool struct
                        tools.push(MCPToolDefinition {
                            service_id,
                            service_name: service_config.name.clone(),
                            name: tool.name.to_string(),
                            description: tool
                                .description
                                .map(|d| d.to_string())
                                .unwrap_or_default(),
                            input_schema: serde_json::Value::Object(
                                tool.input_schema.as_ref().clone(),
                            ),
                        });
                    } else if let Some(name) = tool_value.get("name") {
                        // Handle raw JSON tool definitions
                        let name_str = name.as_str().unwrap_or("unknown").to_string();
                        let description = tool_value
                            .get("description")
                            .and_then(|d| d.as_str())
                            .unwrap_or("")
                            .to_string();
                        let input_schema = tool_value
                            .get("inputSchema")
                            .or_else(|| tool_value.get("input_schema"))
                            .unwrap_or(&serde_json::json!({}))
                            .clone();

                        tools.push(MCPToolDefinition {
                            service_id,
                            service_name: service_config.name.clone(),
                            name: name_str,
                            description,
                            input_schema,
                        });
                    }
                }

                return Ok(tools);
            }
        }

        Ok(vec![])
    }

    // Refresh MCP tools for a service and update data
    pub async fn refresh_mcp_tools(&self, service_id: i64) -> anyhow::Result<()> {
        let tools = self.discover_mcp_tools(service_id).await?;

        // Get service name for the tools
        let _service = self
            .get_mcp_service(service_id)?
            .ok_or_else(|| anyhow::anyhow!("Service not found: {}", service_id))?;

        // Register new tools (this will clear existing tools first)
        self.register_mcp_tools(service_id, &tools)?;

        Ok(())
    }
}

// MCP Service Error Handling
#[derive(Debug, thiserror::Error)]
pub enum MCPServiceError {
    #[error("Service not found: {0}")]
    ServiceNotFound(i64),
    #[error("Service not connected: {0}")]
    ServiceNotConnected(i64),
    #[error("Service connection failed: {0}")]
    ConnectionFailed(String),
    #[error("Service discovery failed: {0}")]
    DiscoveryFailed(String),
    #[error("Service execution failed: {0}")]
    ExecutionFailed(String),
    #[error("Service validation failed: {0}")]
    ValidationFailed(String),
}

impl MCPServiceError {
    pub fn connection_failed<S: Into<String>>(msg: S) -> Self {
        MCPServiceError::ConnectionFailed(msg.into())
    }

    pub fn discovery_failed<S: Into<String>>(msg: S) -> Self {
        MCPServiceError::DiscoveryFailed(msg.into())
    }

    pub fn execution_failed<S: Into<String>>(msg: S) -> Self {
        MCPServiceError::ExecutionFailed(msg.into())
    }

    pub fn validation_failed<S: Into<String>>(msg: S) -> Self {
        MCPServiceError::ValidationFailed(msg.into())
    }
}

// MCP Service Connection Management using rmcp
pub struct MCPServiceManager {
    tool_manager: Arc<ToolManager>,
}

impl MCPServiceManager {
    pub fn new(tool_manager: Arc<ToolManager>) -> Self {
        Self { tool_manager }
    }

    pub async fn connect_to_service(&self, service_id: i64) -> anyhow::Result<()> {
        // Connection is handled by rmcp Service automatically
        // We just need to ensure the service is registered
        let service = self
            .tool_manager
            .get_mcp_service(service_id)?
            .ok_or_else(|| anyhow::anyhow!("Service not found: {}", service_id))?;

        if service.enabled {
            let service_config: ServiceConfig = service.try_into()?;
            let rmcp_service = ServiceWrapper::new(service_config).await?;

            let mut registry = self.tool_manager.service_registry.write().await;
            registry.insert(service_id, rmcp_service);
        }

        Ok(())
    }

    pub async fn disconnect_service(&self, service_id: i64) -> anyhow::Result<()> {
        let mut registry = self.tool_manager.service_registry.write().await;

        // Cleanup stdio transport if it exists
        if let Some(service) = registry.get(&service_id) {
            if let Err(e) = service.cleanup().await {
                eprintln!("Failed to cleanup service {}: {}", service_id, e);
            }
        }

        registry.remove(&service_id);
        Ok(())
    }

    pub async fn reconnect_service(&self, service_id: i64) -> anyhow::Result<()> {
        // Disconnect first
        self.disconnect_service(service_id).await?;

        // Then reconnect
        self.connect_to_service(service_id).await?;

        // Refresh tools after reconnection
        if let Err(e) = self.tool_manager.refresh_mcp_tools(service_id).await {
            eprintln!("Failed to refresh tools after reconnection: {}", e);
        }

        Ok(())
    }

    pub async fn list_active_connections(&self) -> Vec<i64> {
        let registry = self.tool_manager.service_registry.read().await;
        registry.keys().copied().collect()
    }

    pub async fn check_service_health(&self, service_id: i64) -> anyhow::Result<bool> {
        let registry = self.tool_manager.service_registry.read().await;
        if let Some(service) = registry.get(&service_id) {
            // Use rmcp's built-in health check
            match service.is_healthy().await {
                Ok(healthy) => Ok(healthy),
                Err(_) => Ok(false),
            }
        } else {
            Ok(false)
        }
    }

    pub async fn get_service_status(&self, service_id: i64) -> String {
        match self.check_service_health(service_id).await {
            Ok(true) => "healthy".to_string(),
            Ok(false) => "unhealthy".to_string(),
            Err(_) => "error".to_string(),
        }
    }

    pub async fn get_service_info(
        &self,
        service_id: i64,
    ) -> anyhow::Result<Option<MCPServiceConfig>> {
        self.tool_manager.get_mcp_service(service_id)
    }

    pub async fn restart_all_services(&self) -> anyhow::Result<()> {
        let services = self.tool_manager.list_mcp_services()?;
        let mut successful = 0;
        let mut failed = 0;

        for service in services {
            if service.enabled {
                match self.reconnect_service(service.id).await {
                    Ok(_) => successful += 1,
                    Err(e) => {
                        failed += 1;
                        eprintln!("Failed to restart service {}: {}", service.name, e);
                    }
                }
            }
        }

        if failed > 0 {
            eprintln!(
                "Restarted {} services successfully, {} failed",
                successful, failed
            );
        } else {
            eprintln!("Restarted {} services successfully", successful);
        }

        Ok(())
    }
}
